#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Mar 30 16:27:19 2024

@author: lixiao
"""
import numpy as np
 
def normalise2DPts(pts):
    """  normalises 2D homogeneous points

     Function translates and normalises a set of 2D homogeneous points
     so that their centroid is at the origin and their mean distance from
     the origin is sqrt(2).

     Usage:   [pts_tilde, T] = normalise2dpts(pts)

     Argument:
       pts -  3xN array of 2D homogeneous coordinates

     Returns:
       pts_tilde -  3xN array of transformed 2D homogeneous coordinates.
       T         -  The 3x3 transformation matrix, pts_tilde = T*pts
    """
    N = pts.shape[1]

    # Convert homogeneous coordinates to Euclidean coordinates (pixels)
    pts_ = pts/pts[2,:]

    # Centroid (Euclidean coordinates)
    mu = np.mean(pts_[:2,:], axis = 1)

    # Average distance or root mean squared distance of centered points
    # It does not matter too much which criterion to use. Both improve the
    # numerical conditioning of the Fundamental matrix estimation problem.
    pts_centered = (pts_[:2,:].T - mu).T

    # Option 1: RMS distance
    # sigma = np.sqrt( np.mean( np.sum(pts_centered**2, axis = 0) ) )

    # Option 2: average distance
    sigma = np.mean( np.sqrt(sum(pts_centered**2)) )

    s = np.sqrt(2) / sigma
    T = np.array([
        [s, 0, -s * mu[0]],
        [0, s, -s * mu[1]],
        [0, 0, 1]])

    pts_tilde = T @ pts_

    return pts_tilde, T

def fundamentalEightPoint(p1, p2):
    """ The 8-point algorithm for the estimation of the fundamental matrix F

     The eight-point algorithm for the fundamental matrix with a posteriori
     enforcement of the singularity constraint (det(F)=0).
     Does not include data normalization.

     Reference: "Multiple View Geometry" (Hartley & Zisserman 2000), Sect. 10.1 page 262.

     Input: point correspondences
      - p1 np.ndarray(3,N): homogeneous coordinates of 2-D points in image 1
      - p2 np.ndarray(3,N): homogeneous coordinates of 2-D points in image 2

     Output:
      - F np.ndarray(3,3) : fundamental matrix
    """

    # Sanity checks
    assert(p1.shape == p2.shape), "Input points dimension mismatch"
    assert(p1.shape[0] == 3), "Points must have three columns"
    
    num_points = p1.shape[1]
    assert(num_points>=8), \
            'Insufficient number of points to compute fundamental matrix (need >=8)'

    # Compute the measurement matrix A of the linear homogeneous system whose
    # solution is the vector representing the fundamental matrix.
    A = np.zeros((num_points,9))
    for i in range(num_points):
        A[i,:] = np.kron( p1[:,i], p2[:,i] ).T
    
    # "Solve" the linear homogeneous system of equations A*f = 0.
    # The correspondences x1,x2 are exact <=> rank(A)=8 -> there exist an exact solution
    # If measurements are noisy, then rank(A)=9 => there is no exact solution, 
    # seek a least-squares solution.
    _, _, vh= np.linalg.svd(A,full_matrices = True)
    F = np.reshape(vh[-1,:], (3,3)).T

    # Enforce det(F)=0 by projecting F onto the set of 3x3 singular matrices
    u, s, vh = np.linalg.svd(F)
    s[2] = 0
    F = u @ np.diag(s) @ vh

    return F

def fundamentalEightPointNormalized(p1, p2):
    """ Normalized Version of the 8 Point algorith
     Input: point correspondences
      - p1 np.ndarray(3,N): homogeneous coordinates of 2-D points in image 1
      - p2 np.ndarray(3,N): homogeneous coordinates of 2-D points in image 2

     Output:
      - F np.ndarray(3,3) : fundamental matrix
    """

    # Normalize each set of points so that the origin
    # is at centroid and mean distance from origin is sqrt(2).
    p1_tilde, T1 = normalise2DPts(p1)
    p2_tilde, T2 = normalise2DPts(p2)

    # Linear solution
    F = fundamentalEightPoint(p1_tilde, p2_tilde)

    # Undo the normalization
    F = T2.T @ F @ T1

    return F


def fundamentalRANSAC(p1, p2, maxIters=100):
    """
    Takes as input two 3-by-N matrices containing point correspondences between two images
    (in homogeneous coordinates) and computes the fundamental matrix F using a RANSAC routine
    to discard outliers.
    
    Input:
    - p1 np.ndarray(3,N): matrix containing point correspondences in image 1
    - p2 np.ndarray(3,N): matrix containing point correspondences in image 2
    - maxIters int: maximum number of RANSAC iterations (default: 50)
    
    Output:
    - F np.ndarray(3,3): fundamental matrix
    
    """
    # Largest Inlier Set
    inliers = []
    numInliers = -np.Inf
    # Inlier threshold
    inlierThresh = 0.01
    # Main RANSAC loop
    for i in range(maxIters):
        # Randomly sample 8 points
        inds = np.random.choice(range(p1.shape[1]), 8, replace=False)
        # Use these samples to compute the fundamental matrix
        F_cur = fundamentalEightPointNormalized(p1[:, inds], p2[:, inds])
        # Inlier set (current)
        inliers_cur = []
        numInliers_cur = 0
        # Estimate the number of inliers
        for j in range(p1.shape[1]):
            if np.abs(p2[:, j].T @ F_cur @ p1[:, j]) <= inlierThresh:
                inliers_cur.append(j)
                numInliers_cur += 1
        # Best inlier set so far
        if numInliers_cur > numInliers:
            numInliers = numInliers_cur
            inliers = inliers_cur
        # If more than 50 percent of the data are inliers, we assume a
        # reasonable estimate has been found and terminate.
        if numInliers >= 0.5 * p1.shape[1]:
            break
    # Refine the fundamental matrix estimate using the largest inlier set
    F = fundamentalEightPointNormalized(p1[:, inliers], p2[:, inliers])
    
    return F





if __name__ == '__main__':
  p1 = np.loadtxt('./matches0001.txt')
  p2 = np.loadtxt('./matches0002.txt')
  
  p1 = np.r_[p1, np.ones((1, p1.shape[1]))]
  p2 = np.r_[p2, np.ones((1, p2.shape[1]))]
  
  K = np.array([  [1379.74,   0,          760.35],
                  [    0,     1382.08,    503.41],
                  [    0,     0,          1 ]] )
  
  F = fundamentalRANSAC(p1, p2);
  F2 = fundamentalEightPointNormalized(p1, p2)
  print("F:\n", F)
  print('F2:\n', F2)
  print(p1[:,0].T@F2@p2[:,0])






